import warnings
import numpy as np
from mmcv.transforms import to_tensor
from mmcv.transforms.base import BaseTransform
from mmengine.structures import PixelData
from mmseg.registry import TRANSFORMS
from mmseg.structures import SegDataSample


@TRANSFORMS.register_module()
class PackMultiSegInputs(BaseTransform):
    """Pack the inputs data for multi-head semantic segmentation (LCX + LAD).

    输出:
        - inputs: Tensor
        - data_samples: SegDataSample, 包含
            gt_seg_map_lcx / gt_seg_map_lad / gt_sem_seg(合并)
    """

    def __init__(self,
                 meta_keys=('img_path', 'ori_shape', 'img_shape',
                            'pad_shape', 'scale_factor', 'flip',
                            'flip_direction', 'reduce_zero_label')):
        self.meta_keys = meta_keys

    def _pack_single_seg(self, seg_map):
        """转成 PixelData."""
        if len(seg_map.shape) == 2:   # H, W
            data = to_tensor(seg_map[None, ...].astype(np.int64))
        elif len(seg_map.shape) == 3:
            data = to_tensor(seg_map.astype(np.int64))
        else:
            warnings.warn(f'Unexpected seg_map shape {seg_map.shape}')
            data = to_tensor(seg_map.astype(np.int64))
        return PixelData(data=data)

    def transform(self, results: dict) -> dict:
        packed_results = dict()

        # --- 图像 ---
        if 'img' in results:
            img = results['img']
            if len(img.shape) < 3:
                img = np.expand_dims(img, -1)
            img = img.transpose(2, 0, 1)  # HWC -> CHW
            packed_results['inputs'] = to_tensor(img).contiguous()

        # --- DataSample ---
        data_sample = SegDataSample()

        has_lcx, has_lad = False, False
        if 'gt_seg_map_lcx' in results:
            data_sample.gt_seg_map_lcx = self._pack_single_seg(results['gt_seg_map_lcx'])
            has_lcx = True
        if 'gt_seg_map_lad' in results:
            data_sample.gt_seg_map_lad = self._pack_single_seg(results['gt_seg_map_lad'])
            has_lad = True

        # --- 合并成整体血管 mask (lcx ∪ lad) ---
        if has_lcx and has_lad:
            merged = np.logical_or(
                results['gt_seg_map_lcx'] > 0,
                results['gt_seg_map_lad'] > 0
            ).astype(np.uint8)
            data_sample.gt_sem_seg = self._pack_single_seg(merged)

        # --- Meta 信息 ---
        img_meta = {}
        for key in self.meta_keys:
            if key in results:
                img_meta[key] = results[key]
        data_sample.set_metainfo(img_meta)

        packed_results['data_samples'] = data_sample
        return packed_results

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}(meta_keys={self.meta_keys})'
